# from py2neo import Graph
import json
import py2neo
from py2neo import Graph

graph = Graph("bolt://localhost:7687", auth=('neo4j', '00000000'))
subGraph1 = []
subGraph2 = []
subGraph3 = []


def query_graph_old(target):  # 查询以target为头尾节点的三元组
    sub_graph1 = []
    s = "match (head{entity:'" + target + "'})-[r]->(tail) return head, r, tail"
    result = graph.run(s).data()
    for pair in result:
        sub_graph1.append([dict(pair['head'])['entity'], list(pair['r'].types())[0], dict(pair['tail'])['entity']])
    sub_graph2 = []
    result = graph.run("match (head)-[r]->(tail{entity:'" + target + "'}) return head, r, tail").data()
    for pair in result:
        sub_graph2.append([dict(pair['head'])['entity'], list(pair['r'].types())[0], dict(pair['tail'])['entity']])
    return sub_graph1, sub_graph2


def query_relation_old(target):  # 查询以target为关系的三元组
    sub_graph = []
    s = "match (head)-[r:" + target + "]->(tail) return head, r, tail"
    result = graph.run(s).data()
    for pair in result:
        sub_graph.append([dict(pair['head'])['entity'], list(pair['r'].types())[0], dict(pair['tail'])['entity']])
    return sub_graph
def query_relation(target):  # 查询以target为关系的三元组
    sub_graph = []
    s = "match (head)-[r:" + target + "]->(tail) return head, r, tail"
    result = graph.run(s).data()
    for pair in result:
        sub_graph.append([dict(pair['head'])['name'], list(pair['r'].types())[0], dict(pair['tail'])['name']])
    return sub_graph

# 实体 + 关系。
# head
def query_search_old(node_type, node, relation):
    sub_graph1 = []
    que1 = f"match (head:{node_type}" + " {entity:'" + node + "'})-[r:" + relation + "]->(tail) return head, r, tail"
    print(que1)
    result = graph.run(que1).data()
    for pair in result:
        sub_graph1.append([dict(pair['head'])['entity'], list(pair['r'].types())[0], dict(pair['tail'])['entity']])
    sub_graph2 = []
    que2 = "match (head)-[r:" + relation + f"]->(tail:{node_type}" + " {entity:'" + node + "'}) return head, r, tail"
    print(que2)
    result = graph.run(que2).data()
    for pair in result:
        sub_graph2.append([dict(pair['head'])['entity'], list(pair['r'].types())[0], dict(pair['tail'])['entity']])
    subgraph = sub_graph1
    subgraph.extend(sub_graph2)
    return subgraph

def query_search(node_type, node, relation):
    sub_graph1 = []
    que1 = f"match (head:{node_type}" + " {name:'" + node + "'})-[r:" + relation + "]->(tail) return head, r, tail"
    print(que1)
    result = graph.run(que1).data()
    for pair in result:
        sub_graph1.append([dict(pair['head'])['name'], list(pair['r'].types())[0], dict(pair['tail'])['name']])
    sub_graph2 = []
    que2 = "match (head)-[r:" + relation + f"]->(tail:{node_type}" + " {name:'" + node + "'}) return head, r, tail"
    print(que2)
    result = graph.run(que2).data()
    for pair in result:
        sub_graph2.append([dict(pair['head'])['name'], list(pair['r'].types())[0], dict(pair['tail'])['name']])
    subgraph = sub_graph1
    subgraph.extend(sub_graph2)
    return subgraph

# 实体 + 属性
def property_search(node_type, node, property):
    que_property = f"match (d:{node_type}" + " {name:'" + node + "'}) return d." + property + f"AS {property}"
    print(que_property)
    result = graph.run(que_property).data()
    print(result)
    return result

def query_entities_old():
    # 定义 Cypher 查询语句
    query = """
    MATCH (n)
    RETURN DISTINCT labels(n) AS labels, n.entity AS entity
    """
    label_name = ['Disease', 'Class', 'Reason', 'Pathogenesis', 'Symptom', 'Test', 'Test_items', 'Test_Value', 'Drug', 'Frequency', 'Amount', 'Method', 'Treatment', 'Operation', 'ADE', 'Anatomy', 'Level', 'Duration']
    entity_name_list = [[] for i in range(len(label_name))]
    # 执行查询
    results = graph.run(query)
    for record in results:
        labels = record["labels"][0]
        entity_name = record["entity"]
        for j in range(len(label_name)):
            if labels == label_name[j]:
                entity_name_list[j].append(entity_name)
                break
        # print("Labels:", labels, "Entity Name:", entity_name)
    file = open("data/entities.yml", "w", encoding='utf-8')
    file2 = open("data/vocabulary.txt", "w", encoding='utf-8')
    for i in range(len(label_name)):
        file.write(f"- lookup: {label_name[i]}\n")
        file.write(f"  examples: |\n")
        for j in range(len(entity_name_list[i])):
            file.write(f"    - {entity_name_list[i][j]}\n")
            file2.write(f"{entity_name_list[i][j]}\n") # 写入词库
        file.write("\n")
            # print(label_name[i])
        # print(label_name[i])
        # for j in range(len(entity_name_list[i])):
        #     print(entity_name_list[i][j])


# 问题：多张知识图谱中，存在大量重复的节点。唯一性约束没能成功。但也不能不考虑entity_id，因为这是创建关系的依据。
# 解决方案：依次处理每张图（每张图，先加点，再加边）。
# 处理第i张图的实体时，如果entity_type + entity已经存在，则直接修改entity_id为新id，否则创建新点。
# 处理第i张图的关系时，如果先查询关系是否已经存在，如果不存在再创建
def createGraph_old(n):
    for i in range(1, n + 1):
    # for i in range(n, n + 1):
        with open(f'kg/{i}.json', encoding='utf-8') as f:
            jsData = json.loads(f.read())
            # create nodes in i.json
            for paragraph in jsData['paragraphs']:
                for sentence in paragraph['sentences']:
                    for entity in sentence['entities']:
                        # check
                        check = "MATCH (n:" + entity['entity_type'] + " {entity: '" + entity['entity'] + "'}) RETURN COUNT(n) AS count;"
                        # print(check)
                        try:
                            check_res = graph.run(check).data()[0]['count']
                            # print(check_res)
                            # MATCH (n:YourLabel {propertyName: 'oldValue'})
                            # SET n.propertyName = 'newValue'
                            # RETURN n;
                            if check_res > 0: # 仅修改属性
                                modify = "MATCH (n:" + entity['entity_type'] + " {entity: '" + entity['entity'] + "'})" + \
                                    "SET n.entity_id='" + entity['entity_id'] + "'RETURN n"
                                try:
                                    graph.run(modify)
                                except Exception as e:
                                    print(f"error:{e}")
                            else:
                                s = "CREATE(n:" + entity['entity_type'] + "{entity_id:'" + entity[
                                    'entity_id'] + "', entity:'" + \
                                    entity['entity'] + "'})"
                                try:
                                    graph.run(s)
                                except Exception as e:
                                    print(f"error:{e}")
                        except Exception as e:
                            print(f"error:{e}")
            # create relations in i.json
            for paragraph in jsData['paragraphs']:
                for sentence in paragraph['sentences']:
                    for relation in sentence['relations']:
                        # 还需要检查head和tail的节点类型，因为不同类型的节点存在共享id的情况。
                        head_type, tail_type = getType(relation['relation_type'])
                        # print(head_type, end=" ")
                        # print(tail_type)
                        '''
                        MATCH (a:LabelA)-[r:RELATIONSHIP_TYPE]->(b:LabelB)  
                        WHERE a.property = 'valueA' AND b.property = 'valueB'  
                        RETURN EXISTS(r) AS exists;
                        '''
                        check = f"MATCH (a:{head_type})-[r:{relation['relation_type']}]->(b:{tail_type})\
                                WHERE a.entity_id = '{relation['head_entity_id']}' AND b.entity_id = '{relation['tail_entity_id']}'\
                                RETURN COUNT(r) AS count"
                        # print(check)
                        try:
                            check_ans = graph.run(check).data()[0]['count']
                            # print(check_ans)
                            if check_ans == 0:
                                s = f"MATCH (head: {head_type}),(tail:{tail_type}) WHERE head.entity_id = '" + relation[
                                    'head_entity_id'] + "' AND tail.entity_id = '" + relation['tail_entity_id'] + \
                                    "' CREATE (head)-[r:" + relation['relation_type'] + "{relation_id:'" + relation[
                                        'relation_id'] + "'}] -> (tail) RETURN r"
                                try:
                                    graph.run(s)
                                except Exception as e:
                                    print(f"error:{e}")
                        except Exception as e:
                            print(f"error:{e}")
def getType_old(relation_type):
    if "Test_items" in relation_type:
        return "Test_items", "Disease"
    else:
        tmp_types = relation_type.rstrip().split('_')
        return tmp_types[0], tmp_types[1]
def createRelation_old(n):
    for i in range(1, n + 1):
    # for i in range(n, n + 1):
        with open(f'kg/{i}.json', encoding='utf-8') as f:
            jsData = json.loads(f.read())
            for paragraph in jsData['paragraphs']:
                for sentence in paragraph['sentences']:
                    for relation in sentence['relations']:
                        # 还需要检查head和tail的节点类型，因为不同类型的节点存在共享id的情况。
                        head_type, tail_type = getType(relation['relation_type'])
                        # print(head_type, end=" ")
                        # print(tail_type)
                        s = f"MATCH (head: {head_type}),(tail:{tail_type}) WHERE head.entity_id = '" + relation['head_entity_id'] + "' AND tail.entity_id = '" + relation['tail_entity_id'] + \
                            "' CREATE (head)-[r:" + relation['relation_type'] + "{relation_id:'" + relation['relation_id'] + "'}] -> (tail) RETURN r"
                        try:
                            graph.run(s)
                        except:
                            pass


def main():
    # graph.run("match (n) detach delete n")
    # # try:
    # #     graph.run("CREATE CONSTRAINT FOR (n) REQUIRE n.entity IS UNIQUE")  # 唯一性约束
    # # except:
    # #     pass
    # n = 41 # minimum
    # createGraph(n)  # 创建节点
    # # createRelation(n)  # 创建关系
    # query_entities() # 查询实体
    # while True:
    #     target = input()
    #     # re=input()
    #     # subGraph3=query_relation(target)
    #     subGraph1, subGraph2 = query_graph(target)
    #     # subGraph1,subGraph2=query_search(target,re)
    #     print(subGraph1)
    #     print(subGraph2)
        # print(subGraph3)
    return

if __name__ == '__main__':
    main()